import optparse, sys
sys.path.append("/usr/local/bcbb/gff")
from BCBio import GFF
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord


def __main__():
	parser = optparse.OptionParser()
	parser.add_option("-i", "--input", default=None, dest="input", 
					  help="The input list of gene ids")
	parser.add_option("-o", "--output", default=None, dest="gff3_output", 
					  help="The output gff3 file")
	parser.add_option("-g", "--gff3", default=None, dest="gff3",
					  help="The annotated gff3 file")
	(options, args) = parser.parse_args()
	
	#for testing
	limit_info=dict(gff_id=["Chr09"])
	
	#All the genes to search for, these are loci, not transcript ids
	gene_ids = [l.rstrip() for l in open(options.input, "r")]
	
	
	gff_handle = open(options.gff3, "r")
	selected_features = []
	for record in GFF.parse(gff_handle): #, limit_info=limit_info): #all the records on a scaffold/chromosome
		for feature in record.features: #every locus			
			if feature.id in gene_ids:
				seq = Seq("")
				rec = SeqRecord(seq, record.id)
				rec.features.append(feature)
				selected_features.append(rec)
	gff_handle.close()
	
	#save the new gff3 entries
	gff_out_handle = open(options.gff3_output, "w")
	GFF.write(selected_features, gff_out_handle)
	gff_out_handle.close()
	
	
	


if __name__=="__main__": 
	__main__()